import sys
import struct
import zlib
import puza_stream_table as pst
import puza_memory_table as pmt
import puza_row as pr
import logging
import os
import json

class PuzaTableIndex():
    def __init__(self, obj = {}):
        self._local_buffer = obj
        self._bind_table_handler = None
        self._bind_table_file = ''        
    
    def Locate(self, key = '', input_table =''):
        if input_table != '':
            ## reset bind table
            self._bind_table_file = input_table
            self._bind_table_handler = open(self._bind_table_file, 'rb')              
        else:
            if self._bind_table_file == None:
                return None
        records = self._local_buffer[key]
        results = []
        for record in records:
            block_number = record[0]
            read_length = record[1]
            record_number = record[2]
            self._bind_table_handler.seek(block_number)
            s = self._bind_table_handler.read(read_length)
            decom_s = zlib.decompress(s)
            w = decom_s.strip().split('\n')
            r = pr.PuzaRow.Load(w[record_number])
            results.append(r)
        return results

class PuzaTableIndexer:
    def __init__(self):
        pass

    @staticmethod
    def MakeIndex(input_filename = '', output_filename = '', key = '', lformat=0):
        if input_filename != '' and key != '' and output_filename != '':
            input_file = open(input_filename, 'rb')
            c = input_file.read(4) ## skip the magic number
            local_index= {}
            N = 0
            if pst.PuzaStreamTable.CheckMagic(c):
                c = input_file.read(4)
                while c != "":
                    l = struct.unpack('i',c)[0] ## read the block size
                    block_byte = input_file.tell()
                    s = input_file.read(l)
                    decom_s = zlib.decompress(s)
                    w = decom_s.strip().split('\n')
                    for i in range(len(w)):
                        r = pr.PuzaRow.Load(w[i])
                        target_key_value = r[key]
                        if target_key_value != None:
                            record_offset = i
                            if not target_key_value in local_index:
                                local_index[target_key_value] = []
                            local_index[target_key_value].append((block_byte, l, record_offset))
                            N = N + 1
                    c = input_file.read(4)
            input_file.close()
            ## save index
            output_table = pst.PuzaStreamTable()
            output_table.OutputFile(output_filename)
            for local_key, value in local_index.iteritems():
                r = pr.PuzaRow()
                r[str(local_key)] = value
                output_table.StreamOutRow(r, lformat=lformat)
            r = pr.PuzaRow()
            r['key'] = key
            r['key_counts'] = len(local_index)
            r['item_counts'] = N
            logging.info("Key Counts:" + str(len(local_index)))
            logging.info("Item Counts:" + str(N))
            output_table.StreamOutRow(r, lformat=lformat)
            output_table.StreamClose()
    
    @staticmethod
    def MakeIndexStream(input_filename = '', output_filename = '', key = '', lformat=0, block_number = 256):
        step1_file = output_filename + '.temp'
        logging.info("Starting Step1")
        PuzaTableIndexer.MakeIndexStreamStep1(input_filename, step1_file, key, lformat)
        logging.info("Starting Step2")
        local_buffer = PuzaTableIndexer.MakeIndexStreamStep2(step1_file)
        logging.info("Starting Step3")
        PuzaTableIndexer.MakeIndexStreamStep3(local_buffer, output_filename, key, lformat, block_number)
        #os.remove(step1_file)        
    
    @staticmethod
    def MakeIndexStreamStep3(local_buffer, output_filename, key, lformat, block_number):
        last_key_value = ''
        small_buffer = []
        OutputTable = pst.PuzaStreamTable()
        OutputTable.OutputFile(output_filename)
        K = 0
        N = 0
        for i in range(len(local_buffer)):
            key_value = local_buffer[i][0]
            if key_value != last_key_value and last_key_value != '':
                ## output
                r = pr.PuzaRow()
                r[str(last_key_value)] = small_buffer
                OutputTable.StreamOutRow(r, lformat, block_number)
                K = K + 1
                N = N + len(small_buffer)
                small_buffer = []
            last_key_value = key_value
            small_buffer.append((local_buffer[i][1], local_buffer[i][2], local_buffer[i][3]))
        ## check remain
        if len(small_buffer) > 0:
            ## output
            r = pr.PuzaRow()
            r[str(last_key_value)] = small_buffer
            OutputTable.StreamOutRow(r, lformat, block_number)
            K = K + 1
            N = N + len(small_buffer)
        r = pr.PuzaRow()
        r['key'] = key
        r['key_counts'] = K
        r['item_counts'] = N
        OutputTable.StreamOutRow(r, lformat, block_number)
        logging.info("Key Counts:" + str(K))
        logging.info("Item Counts:" + str(N))
        OutputTable.StreamClose()
        
    @staticmethod
    def MakeIndexStreamStep2(input_filename = ''):
        InputTable = pst.PuzaStreamTable()
        InputTable.InputFile(input_filename)
        local_buffer = []
        N = 0
        while True:
            raw_obj = InputTable.StreamInRow()
            if raw_obj != None:
                r = pr.PuzaRow.Load(raw_obj)
                key_value = r['i'][0]
                block_byte = r['i'][1]
                l = r['i'][2]
                offset = r['i'][3]
                local_buffer.append((key_value, block_byte, l, offset))
                N = N + 1
                if N % 1000000 == 0:
                    logging.info('Processed ' + str(N))
            else:
                break
        local_buffer = sorted(local_buffer, key = lambda x:x[0])
        return local_buffer
            
        
    @staticmethod
    def MakeIndexStreamStep1(input_filename = '', output_filename = '', key = '', lformat=0):
        if input_filename != '' and key != '' and output_filename != '':
            input_file = open(input_filename, 'rb')
            c = input_file.read(4) ## skip the magic number
            N = 0
            output_table = pst.PuzaStreamTable()
            output_table.OutputFile(output_filename)
            if pst.PuzaStreamTable.CheckMagic(c):
                c = input_file.read(4)
                while c != "":
                    l = struct.unpack('i',c)[0] ## read the block size
                    block_byte = input_file.tell()
                    s = input_file.read(l)
                    decom_s = zlib.decompress(s)
                    w = decom_s.strip().split('\n')
                    for i in range(len(w)):
                        r = pr.PuzaRow.Load(w[i])
                        target_key_value = r[key]
                        if target_key_value != None:
                            record_offset = i
                            r = pr.PuzaRow()
                            r['i'] = ((str(target_key_value), block_byte, l, record_offset))
                            output_table.StreamOutRow(r, lformat=lformat)
                            N = N + 1
                            if N % 1000000 == 0:
                                logging.info('Processed ' + str(N))
                    c = input_file.read(4)
            input_file.close()            
            output_table.StreamClose()

    @staticmethod
    def ReadIndex(input_filename = ''):
        if input_filename != '':
            input_table = pst.PuzaStreamTable()
            input_table.InputFile(input_filename)
            local_buffer = {}
            while True:
                raw_obj = input_table.StreamInRow()
                if raw_obj != None:
                    obj = json.loads(raw_obj)
                    key_name = obj.keys()
                    for key_name, key_value in obj.iteritems():
                        local_buffer[key_name] = key_value
                else:
                    break
            pindex = PuzaTableIndex(local_buffer)
            return pindex
